#!/usr/bin/env bash
# Copyright (c) 2024 maminjie <canpool@163.com>
# SPDX-License-Identifier: MulanPSL-2.0

ROOT_DIR=$(dirname $(readlink -f "$0"))
CUR_DIR=$(pwd)
WORK_DIR=$CUR_DIR/iso_cut

MOUNT_DIR=$WORK_DIR/mount
ISO_DIR=$WORK_DIR/iso
YUM_DIR="$WORK_DIR/yum"
YUM_CONFIG_FILE="$YUM_DIR/yum.conf"

ARCH="$(uname -m)"
DATE="$(date +%Y%m%d)"
ISO_INFILE=""
ISO_OUTFILE=""
REPOS=""
PKG_FILES=""
EXT_FILE=""
KS_FILE=""
HEADLESS_FLAG=0


usage() {
printf "cut operating system iso image

usage:
    bash iso_cut.sh [-h] [-i OLDISO] [-o NEWISO] [-s REPOSITORY] [-f FILE] [-e EXTFILE] [-k KSFILE]

optional arguments:
    -i OLDISO      openeuler iso filepath
    -o NEWISO      new iso filepath, default is new-ARCH-dvd-YYYYMMDD.iso in current dir
    -s REPOSITORY  source dnf repository address link (can be specified multiple times)
    -f FILE        package list file, contains extra packages (can be specified multiple times)
    -e EXTFILE     extra package list file, not installed by default
    -k KSFILE      kickstart filepath
    -a             perform a headless installation
    -h             show the help message and exit
\n"
}


log_error() {
    echo "ERROR: $1"
}

log_info() {
    echo "INFO: $1"
}


parse_cmdline() {
    if [ $# -eq 0 ]; then
        usage; exit 0
    fi

    if [[ "${ARCH}" != "aarch64" && "${ARCH}" != "x86_64" ]]; then
        log_error "unsupported architecture: ${ARCH}"; return 1
    fi

    # parse cmd line arguments
    while getopts ":i:o:s:f:e:k:ah" opt; do
        case "$opt" in
            i)
                ISO_INFILE="$OPTARG"
            ;;
            o)
                ISO_OUTFILE="$OPTARG"
            ;;
            s)
                if [ -z "$REPOS" ]; then
                    REPOS="$OPTARG"
                else
                    REPOS="$REPOS $OPTARG"
                fi
            ;;
            f)
                if [ -f "$OPTARG" ]; then
                    if [ -z "$PKG_FILES" ]; then
                        PKG_FILES="$OPTARG"
                    else
                        PKG_FILES="$PKG_FILES $OPTARG"
                    fi
                else
                    log_error "$OPTARG is nof file"; return 1
                fi
            ;;
            e)
                EXT_FILE="$OPTARG"
            ;;
            k)
                KS_FILE="$OPTARG"
            ;;
            a)
                HEADLESS_FLAG=1
            ;;
            h)
                usage; exit 0
            ;;
            ?)
                echo "please check the params."; usage; return 1
            ;;
        esac
    done
    if [ ! -f "$ISO_INFILE" ]; then
        echo "invalid ISO arguments"; return 1
    fi
    if [ -z "$ISO_OUTFILE" ]; then
        ISO_OUTFILE="new-$ARCH-dvd-$DATE.iso"
    fi
    ISO_OUTFILE=$(realpath $ISO_OUTFILE)
    return 0
}

check_user() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root!"
        exit 1
    fi
}

check_requires() {
    local ret=0

    # createrepo_c package contains:
    #   createrepo
    # dnf-plugins-core package contains:
    #   repoclosure yumdownloader
    # genisoimage package contains:
    #   mkisofs
    # isomd5sum package contains
    #   implantisomd5
    # eg:
    #   dnf install createrepo_c dnf-plugins-core genisoimage isomd5sum -y
    for c in yumdownloader mkisofs createrepo implantisomd5; do
        command -v $c > /dev/null
        if [ $? -ne 0 ]; then
            log_error "command $c not found"
            ret=1
        fi
    done

    return $ret
}

# make_iso iso_label iso_file iso_dir
# params:
#   iso_label - equal to LABEL in EFI/BOOT/grub.cfg
#   iso_file - new iso file name
#   iso_dir - direcotry of iso files
make_iso() {
    if [ $# -lt 3 ]; then
        echo "invalid param"; return 1
    fi

    local arch=$(uname -m)
    local iso_label=$1
    local iso_file=$2
    local iso_dir=$(realpath $3)

    if [ ! -d "$iso_dir" ]; then
        echo "$iso_dir: no such directory"; return 1
    fi

    if [[ "$arch" = "x86_64" ]]; then
        mkisofs -R -J -T -r -l -d -joliet-long -allow-multidot -allow-leading-dots -no-bak -V "${iso_label}" -o "${iso_file}" -b isolinux/isolinux.bin \
            -c isolinux/boot.cat -no-emul-boot -boot-load-size 4 -boot-info-table  -eltorito-alt-boot -e images/efiboot.img -no-emul-boot "${iso_dir}"
        [[ $? -ne 0 ]] && return 1
        isohybrid -u "$iso_file"
    elif [[ "$arch" = "aarch64" || "$arch" = "loongarch64" ]]; then
        mkisofs -R -J -T -r -l -d -joliet-long -allow-multidot -allow-leading-dots -no-bak -V "${iso_label}" -o "${iso_file}" -e images/efiboot.img -no-emul-boot "${iso_dir}"
        [[ $? -ne 0 ]] && return 1
    else
        echo "unsupported architecture: ${arch}"; return 1
    fi
    implantisomd5 "$iso_file"
    return 0
}

# create_yum_repo YUM_DIR YUM_CONFIG_FILE REPOS
create_yum_repo() {
    local yum_dir="$1"
    local yum_config_file="$2"
    local repos="$3"
    local yum_repo_dir="${yum_dir}/yum.repos.d"
    local yum_repo_file="${yum_repo_dir}/os.repo"

    mkdir -p "$yum_repo_dir"

    cat > "${yum_config_file}" <<EOF
[main]
gpgcheck=1
installonly_limit=3
clean_requirements_on_remove=True
best=True
skip_if_unavailable=False
reposdir=${yum_repo_dir}
EOF

    local i=1
    for repo in $repos; do
        cat >> "${yum_repo_file}" <<EOF
[os_repo_$i]
name=$repo
baseurl=$repo
enabled=1
gpgcheck=0

EOF
        ((i++))
    done
}

init_dir() {
    [[ ! -d "$1" ]] && mkdir -p "$1"
}

clean_dir() {
    [[ -d "$1" ]] && rm -rf "$1"
}

cut_iso() {
    # init dirs
    init_dir $WORK_DIR
    init_dir $MOUNT_DIR

    # mount iso
    umount $MOUNT_DIR 2> /dev/null
    mount $ISO_INFILE $MOUNT_DIR > /dev/null

    # config repo
    init_dir $YUM_DIR
    local yum_tmp_dir="$YUM_DIR/tmp"
    mkdir -p $yum_tmp_dir
    REPOS="file://$MOUNT_DIR $REPOS"
    create_yum_repo "$YUM_DIR" "$YUM_CONFIG_FILE" "$REPOS"

    # make iso dir
    log_info "make iso dir..."
    clean_dir $ISO_DIR
    init_dir $ISO_DIR
    log_info "prepare install image..."
    for i in $(ls $MOUNT_DIR); do
        if [[ "$i" = "Packages" || "$i" = "repodata" ]]; then
            continue
        fi
        cp -ar $MOUNT_DIR/$i $ISO_DIR
    done
    mkdir -p $ISO_DIR/{Packages,repodata}
    log_info "prepare packages..."
    if [ -z "$PKG_FILES" ]; then
        cp -ar $MOUNT_DIR/Packages $ISO_DIR
        if [ -f "$MOUNT_DIR/repodata/normal.xml" ]; then
            cp -a $MOUNT_DIR/repodata/normal.xml $ISO_DIR/repodata/normal.xml
        fi
    else
        cp $ROOT_DIR/config/repodata.template $ISO_DIR/repodata/normal.xml
        for f in $PKG_FILES; do
            for pkg in $(cat $f); do
                sed -i "/package_here/a <packagereq type=\"default\">$pkg<\/packagereq>" $ISO_DIR/repodata/normal.xml
            done
            yumdownloader -y --resolve -c "$YUM_CONFIG_FILE" --installroot="$yum_tmp_dir" --destdir="$ISO_DIR/Packages/" $(tr '\n' ' ' < "$f")
            if [ $? -ne 0 ]; then
                log_error "yumdownloader packages from $f failed"; return 1
            fi
        done
    fi
    if [ -f "$EXT_FILE" ]; then
        local f=$EXT_FILE
        yumdownloader -y --resolve -c "$YUM_CONFIG_FILE" --installroot="$yum_tmp_dir" --destdir="$ISO_DIR/Packages/" $(tr '\n' ' ' < "$f")
        if [ $? -ne 0 ]; then
            log_error "yumdownloader packages from $f failed"; return 1
        fi
    fi
    log_info "create repo..."
    if [ -f "$ISO_DIR/repodata/normal.xml" ]; then
        createrepo -g $ISO_DIR/repodata/normal.xml $ISO_DIR
    elif [ -f "$ISO_DIR/comps.xml" ]; then
        # for BCLinux-for-Euler
        createrepo -g $ISO_DIR/comps.xml $ISO_DIR
    else
        createrepo $ISO_DIR
    fi

    local isolinux_cfg="$ISO_DIR/isolinux/isolinux.cfg"
    local efilinux_cfg="$ISO_DIR/EFI/BOOT/grub.cfg"
    # config kickstart
    if [ -f "$KS_FILE" ]; then
        local ks_cfg="custom-ks.cfg"
        cp $KS_FILE $ISO_DIR/$ks_cfg
        if [ -f "$isolinux_cfg" ]; then
            sed -i "/append/ s/$/ inst.ks=cdrom:\/$ks_cfg/g" $isolinux_cfg
        fi
        sed -i "/inst.stage2/ s/$/ inst.ks=cdrom:\/$ks_cfg/g" $efilinux_cfg
    fi

    # config headless installation
    if [ $HEADLESS_FLAG -eq 1 ]; then
        if [ -f "$isolinux_cfg" ]; then
            # NOTE: timeout cannot be set to 0, and the menu default must be set
            sed -i 's/^timeout.*/timeout 10/g' $isolinux_cfg
            grep -q "menu label ^Install" $isolinux_cfg
            if [ $? -eq 0 ]; then
                sed -i '/menu default/d' $isolinux_cfg
                sed -i '/menu label \^Install/a \ \ menu default' $isolinux_cfg
            fi
        fi
        sed -i 's/^set timeout=.*/set timeout=0/g' $efilinux_cfg
        sed -i 's/^set default=.*/set default="0"/g' $efilinux_cfg
    fi

    # remake iso
    log_info "remake iso image ..."
    local iso_label=$(grep -roE "LABEL=[^ ]*" $efilinux_cfg | head -1 | cut -d '=' -f 2)
    make_iso $iso_label $ISO_OUTFILE $ISO_DIR
    if [ $? -ne 0 ]; then
        log_error "make iso failed"
        return 1
    fi

    # clean
    umount $MOUNT_DIR 2> /dev/null
    cd $CUR_DIR
    clean_dir $WORK_DIR

    return 0
}

main() {
    check_user

    parse_cmdline "$@"
    if [ $? -ne 0 ]; then
        log_error "parse params failed, try -h for help"
        return 1
    fi

    check_requires
    if [ $? -ne 0 ]; then
        log_error "check requires failed"
        return 1
    fi

    cut_iso
    if [ $? -ne 0 ]; then
        log_error "cut iso failed"
        return 1
    fi

    return 0
}

main "$@"
